#!/usr/bin/env python
import rospy
from franka_msgs.msg import FrankaState
from franka_example_controllers.msg import LhyCartesianCmd
from franka_example_controllers.msg import LhyCmdState
import numpy as np
import time


class FrankaPosServo:
    """
    A simple class to control Franka robot in Cartesian space.
    It can send commands to the robot in both blocking and non-blocking ways.
    """
    def __init__(self):
        rospy.init_node('grasper',anonymous=True)
        time.sleep(0.5)
        self.pub = rospy.Publisher('/cartesian_position_controller/target',LhyCartesianCmd,queue_size=1,latch=True)

        self.robot_sub = rospy.Subscriber('/franka_state_controller/franka_states',FrankaState,self._get_robot_position)
        self.control_sub = rospy.Subscriber('/cartesian_position_controller/is_running',LhyCmdState, self._get_robot_running)
        

        self.robot_position = None # The current position of the robot in Cartesian space, 3d array-like.
        self.robot_is_running = False # Whether the robot is currently running a command.

        self.send_cmd_id = None
        self.recv_cmd_id = None
    
        while self.robot_position is None:
            rospy.sleep(0.1)
        while self.recv_cmd_id is None:
            rospy.sleep(0.1)
        self.send_cmd_id = (self.recv_cmd_id + 1) % (1<<31)
        print("Franka Servo initialized successfully!")
    
    def gen_abs_move_cmd(self, goal_position, max_speed: float = 0.01):
        """
        Generate a command to move the robot to a specific position in Cartesian space.
        Args:
          goal_position: The target position in Cartesian space, should be a 3d array-like.
          max_speed: The maximum speed of the robot, default is 0.01 m/s.
        Returns:
            A LhyCartesianCmd object that can be published to the robot.
        """

        goal_position, max_speed = self._type_check(goal_position, max_speed)
        return self._get_new_cmd(goal_position,max_speed)
    
    def gen_rela_move_cmd(self, relative_move, max_speed: float = 0.01):
        """
        Generate a command to move the robot relative to its current position in Cartesian space.
        Args:
          relative_move: The relative movement in Cartesian space, should be a 3d array-like.
          max_speed: The maximum speed of the robot, default is 0.01 m/s.
        Returns:
            A LhyCartesianCmd object that can be published to the robot.
        """
        relative_move, max_speed = self._type_check(relative_move, max_speed)
        return self._get_new_cmd(self.robot_position + relative_move, max_speed)
    
    def non_blocking_move(self, cmd: LhyCartesianCmd):
        """
        Send a command to the robot in a non-blocking way.
        The function will return immediately after publishing the command.
        Args:
            cmd: A LhyCartesianCmd object that can be published to the robot.
        """
        self.pub.publish(cmd)
    
    def blocking_move(self, cmd: LhyCartesianCmd):
        """
        Send a command to the robot in a blocking way.
        The function will block until the robot reaches the target position.
        Args:
            cmd: A LhyCartesianCmd object that can be published to the robot.
        """

        cmd_id = cmd.cmd_id
        self.pub.publish(cmd)
        while not rospy.is_shutdown():
            if self.is_cmd_finished(cmd_id):
                print("Reached goal position.")
                break
            rospy.sleep(0.2)
            if cmd_id != self.recv_cmd_id:
                cmd = self._restart_cmd(cmd)
                self.pub.publish(cmd)
                cmd_id = cmd.cmd_id
                print("Warning: A command is lost.")
    
    def is_cmd_finished(self, cmd_id: int) -> bool:
        """
        Check if a command is finished.
        Args:
            cmd_id: The command ID to check.
        Returns:
            True if the command is finished, False otherwise.
        """
        return cmd_id == self.recv_cmd_id and not self.robot_is_running

    
    def _get_robot_position(self,state: FrankaState):
        otee = state.O_T_EE
        self.robot_position = np.array(otee[12:15])
    
    def _get_robot_running(self, state: LhyCmdState):
        self.robot_is_running = state.is_running
        self.recv_cmd_id = state.cmd_id

    def _get_new_cmd(self, goal_position: np.ndarray, max_speed: float) -> LhyCartesianCmd:
        new_cmd = LhyCartesianCmd()
        new_cmd.cmd_id = self.send_cmd_id
        new_cmd.position_c = goal_position.tolist()
        new_cmd.max_speed_c = max_speed

        # self.pub.publish(new_cmd)
        self.send_cmd_id += 1
        if self.send_cmd_id >= (1<<31):
            self.send_cmd_id = 0
        return new_cmd 
    
    def _restart_cmd(self, cmd: LhyCartesianCmd) -> LhyCartesianCmd:
        cmd.cmd_id = self.send_cmd_id
        self.send_cmd_id += 1
        if self.send_cmd_id >= (1<<31):
            self.send_cmd_id = 0
        return cmd

    
    def _type_check(self, goal_position, max_speed):
        if type(goal_position) == list:
            assert len(goal_position) == 3, "goal_position should be a 3d array-like."
            goal_position = np.array(goal_position)
        elif type(goal_position) == np.ndarray:
            assert goal_position.size == 3, "goal_position should be a 3d array-like."
            goal_position.squeeze()
        else:
            raise TypeError("goal position should be a list or a ndarray.")
    
        return goal_position, max_speed
            

if __name__ == "__main__":
    franka_servo = FrankaPosServo()
    original_position = franka_servo.robot_position
    cmd = franka_servo.gen_rela_move_cmd([0,0,0.05])
    franka_servo.non_blocking_move(cmd)
    time.sleep(3)
    
    new_cmd = franka_servo.gen_abs_move_cmd(original_position)
    franka_servo.blocking_move(new_cmd)

